from parse.stmt import *
from parse.expr import *

from typing import List
from enum import IntEnum


class Resolver(StmtVisitor, ExprVisitor):
    """
    闭包变量的捕捉
    """

    def __init__(self, stmts: List[Stmt]):
        self.stmts = stmts
        self.scopes: List['Scope'] = []

    def visit_binary(self, expr: 'ExprBinary'):
        pass

    def visit_grouping(self, expr: 'ExprGrouping'):
        expr.expr.accept(self)

    def visit_literal(self, expr: 'ExprLiteral'):
        pass

    def visit_unary(self, expr: 'ExprUnary'):
        pass

    def visit_variable(self, expr: 'ExprVariable'):
        if not self.resolve_local(expr.name.lexeme):
            if not self.resolve_upper(expr.name.lexeme):
                raise ValueError(f'undefined variable {expr.name.lexeme}')

    def visit_assign(self, expr: 'ExprAssign'):
        expr.expr.accept(self)
        if not self.resolve_local(expr.name.lexeme):
            if not self.resolve_upper(expr.name.lexeme):
                raise ValueError(f'undefined variable {expr.name.lexeme}')
        # self.resolve_upper(expr.)

    def visit_logical(self, expr: 'ExprLogical'):
        pass

    def visit_call(self, expr: 'ExprCall'):
        pass

    def visit_get(self, expr: 'ExprGet'):
        pass

    def visit_set(self, expr: 'ExprSet'):
        pass

    def visit_this(self, expr: 'This'):
        pass

    def visit_super(self, expr: 'Super'):
        pass

    def visit_print_stmt(self, stmt: 'StmtPrint'):
        pass

    def visit_expr_stmt(self, stmt: 'StmtExpr'):
        stmt.expr.accept(self)

    def visit_var_stmt(self, stmt: 'StmtVar'):
        if stmt.name.lexeme in self.scopes[-1]:
            raise ValueError(f'duplicated variable {stmt.name.lexeme}')

        if len(self.scopes) == 1:
            var_type = VariableType.GLOBAL
        else:
            var_type = VariableType.LOCAL

        self.scopes[-1].add_variable(stmt.name.lexeme, stmt, var_type)

    def visit_block_stmt(self, stmt: 'StmtBlock'):
        self.begin_scope(scope_type=ScopeType.PLAIN)
        for s in stmt.stmts:
            s.accept(self)

        self.end_scope()

    def visit_if_stmt(self, stmt: 'StmtIf'):
        stmt.then_stmt.accept(self)
        if stmt.else_stmt is not None:
            stmt.else_stmt.accept(self)

    def visit_while_stmt(self, stmt: 'StmtWhile'):
        stmt.body.accept(self)

    def visit_break_stmt(self, stmt: 'StmtBreak'):
        pass

    def visit_continue_stmt(self):
        pass

    def visit_class_stmt(self, stmt: 'StmtClass'):
        self.begin_scope(ScopeType.CLASS)
        for s in stmt.methods:
            s.accept(self)

        self.end_scope()

    def visit_function_stmt(self, stmt: 'StmtFun'):
        scope_type = FunType.FUN if stmt.fun_type == FunType.FUN else FunType.METHOD
        self.begin_scope(scope_type)
        param_names = [p.lexeme for p in stmt.params]
        for i, p in enumerate(param_names):
            if param_names.count(p) != 1:
                raise ValueError(f'duplicated param {p}')

            self.scopes[-1].add_variable(p, StmtVar(stmt.params[i], ExprLiteral(None)))

        for s in stmt.body:
            s.accept(self)

        self.end_scope()

    def visit_return_stmt(self, stmt: 'StmtReturn'):
        pass

    def resolve(self):
        self.begin_scope(ScopeType.PLAIN)
        for stmt in self.stmts:
            # print(f'resolve statement {stmt}')
            stmt.accept(self)

        self.end_scope()

    def resolve_local(self, v: str) -> Optional['Variable']:
        for var in self.scopes[-1].variables:
            if var.name == v:
                return var

        return None

    def resolve_upper(self, v: str) -> Optional['Variable']:
        """
        解析变量
        若在某个scope中找到该变量，则设置其is_free属性，表明其将被嵌套定义的函数捕获，需要分配在堆上，并且
        对于后续的scope，需要添加该变量到其调用栈上
        todo (优化): 目前情况下，只要一个变量是自由变量，就必须将其传递到所有嵌套定义的函数内；可以在增加相应的属性，在编译时就判断该变量是否在需要在编译时被某个闭包所捕获
        """
        ret = None
        scope_depth = -1
        for scope in reversed(self.scopes[:-1]):
            for var in scope.variables:
                if var.name == v:
                    if var.var_type == VariableType.LOCAL:
                        var.var_type = VariableType.UPVALUE
                        var.stmt.local = False
                    ret = var
                    scope_depth = scope.depth
                    break

        if ret is not None:
            for scope in self.scopes[scope_depth:]:
                scope.add_variable(v, None, VariableType.UPVALUE)

        return ret

    def begin_scope(self, scope_type):
        self.scopes.append(Scope(len(self.scopes), scope_type))

    def end_scope(self):
        self.scopes.pop()


class ScopeType(Enum):
    PLAIN = auto()
    FUNC = auto()
    LOOP = auto()
    MAIN = auto()
    CLASS = auto()
    METHOD = auto()


class Scope:
    def __init__(self, depth: int, scope_type: 'ScopeType'):
        self.depth = depth
        self.variables: List['Variable'] = []
        self.loop_break_points = []
        self.loop_continue_points = []
        self.scope_type = scope_type
        self.loop_start = 0

    def add_variable(self, name: str, stmt: Optional[StmtVar] = None, var_type=None):
        t = VariableType.LOCAL if var_type is None else var_type
        v = Variable(name, len(self.variables), t, stmt)
        self.variables.append(v)

        return v

    def __contains__(self, item):
        for v in self.variables:
            if v.name == v:
                return True

        return False

    def resolve_variable(self, v: str) -> Optional['Variable']:
        for var in self.variables:
            if var.name == v:
                return var

        return None


class Variable:
    def __init__(self, name: str, offset: int, var_type: 'VariableType', stmt: Optional[StmtVar] = None):
        self.name = name
        self.stmt = stmt
        self.offset = offset
        self.var_type = var_type


class VariableType(IntEnum):
    GLOBAL = auto()
    UPVALUE = auto()
    LOCAL = auto()
    FREE = auto()  # 自由变量，供嵌套定义的函数/方法捕捉作为UPVALUE；声明时需要将其分配到堆上
